import argparse
from collections import defaultdict

def read_bed(input_file):
    sites=defaultdict(list)
    with open(input_file, 'r') as rf:
        for line in rf:
            words = line.strip().split("\t")
            sites[(words[0],words[5])].append(line)
    return sites

def filter(sites,outfilefh1,outfilefh2,outfilefh3,kmer=21):
    kmer=kmer//2
    for key in sites.keys():
        for i in range(len(sites[key])):
            iline=sites[key][i]
            iwords=iline.strip().split("\t")
            ipos=iwords[1]
            left_flag=0
            right_flag=0
            for j in range(max(0,i-kmer),min(i+kmer,len(sites[key]))):
                jline=sites[key][j]
                jwords=jline.strip().split("\t")
                jpos=jwords[1]
                #if int(jpos)==int(ipos):
                #    continue
                if int(jpos)<int(ipos) and int(jpos)>int(ipos)-kmer:
                    left_flag=1
                elif int(jpos)>int(ipos) and int(jpos)<int(ipos)+kmer:
                    right_flag=1
            if left_flag==1 and right_flag!=1:
                outfilefh1.write(iline)
            elif left_flag!=1 and right_flag==1:
                outfilefh2.write(iline)
            else:
                outfilefh3.write(iline)
    

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input",
        type=str,
        default="/home/xiaoyifu/data/HG002/Bisulfite/non_singleton.bisulfite.bed",
        help="input address.",
    )
    parser.add_argument(
        "--ns1",
        type=str,
        default="/home/xiaoyifu/data/HG002/Bisulfite/non_singleton.bisulfite.left",
        help="result address.",
    )
    parser.add_argument(
        "--ns2",
        type=str,
        default="/home/xiaoyifu/data/HG002/Bisulfite/non_singleton.bisulfite.right",
        help="result address.",
    )
    parser.add_argument(
        "--ns3",
        type=str,
        default="/home/xiaoyifu/data/HG002/Bisulfite/non_singleton.bisulfite.both",
        help="result address.",
    )
    parser.add_argument(
        "--kmer",
        type=int,
        default=21,
        help="kmer.",
    )
    return parser.parse_args()

if __name__ == '__main__':
    args=parse_args()
    sites=read_bed(args.input)
    nonsingletonfh1 = open(args.ns1+'.kmer.'+str(args.kmer)+'.bed', "w", buffering=512)
    nonsingletonfh2 = open(args.ns2+'.kmer.'+str(args.kmer)+'.bed', "w", buffering=512)
    nonsingletonfh3 = open(args.ns3+'.kmer.'+str(args.kmer)+'.bed', "w", buffering=512)
    filter(sites,nonsingletonfh1,nonsingletonfh2 ,nonsingletonfh3,args.kmer)